"""
Visualization utilities for attribution analysis results.
"""
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
from typing import List, Dict, Any, Optional
from ..models.schemas import DrillDownResult, MetricDecompositionResult, AnomalyResult


class Visualizer:
    """Creates visualizations for attribution analysis."""

    def __init__(self):
        self.default_template = "plotly_white"

    def create_drill_down_chart(
        self,
        drill_result: DrillDownResult,
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a bar chart for drill-down results.

        Args:
            drill_result: DrillDownResult to visualize
            title: Optional chart title

        Returns:
            Plotly figure
        """
        if not drill_result.top_contributors:
            return go.Figure()

        # Prepare data
        labels = [c.value for c in drill_result.top_contributors]
        values = [c.metric_value for c in drill_result.top_contributors]
        contributions = [c.contribution_percentage for c in drill_result.top_contributors]

        # Create figure
        fig = go.Figure()

        fig.add_trace(go.Bar(
            x=labels,
            y=values,
            text=[f"{c:.1f}%" for c in contributions],
            textposition='auto',
            marker_color='steelblue',
            hovertemplate='<b>%{x}</b><br>Value: %{y:,.0f}<br>Contribution: %{text}<extra></extra>'
        ))

        chart_title = title or f"Top Contributors - {drill_result.dimension} (Level {drill_result.level})"

        fig.update_layout(
            title=chart_title,
            xaxis_title=drill_result.dimension,
            yaxis_title="Metric Value",
            template=self.default_template,
            height=400
        )

        return fig

    def create_waterfall_chart(
        self,
        decomposition: MetricDecompositionResult,
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a waterfall chart for metric decomposition.

        Args:
            decomposition: MetricDecompositionResult to visualize
            title: Optional chart title

        Returns:
            Plotly figure
        """
        # Prepare data
        components = list(decomposition.component_changes.keys())
        changes = [decomposition.component_changes[c] for c in components]

        # Create waterfall chart
        fig = go.Figure(go.Waterfall(
            name="Decomposition",
            orientation="v",
            measure=["relative"] * len(components) + ["total"],
            x=components + ["Total Change"],
            y=changes + [sum(changes)],
            text=[f"{v:,.0f}" for v in changes + [sum(changes)]],
            textposition="outside",
            connector={"line": {"color": "rgb(63, 63, 63)"}},
        ))

        chart_title = title or f"Metric Decomposition - {decomposition.metric}"

        fig.update_layout(
            title=chart_title,
            template=self.default_template,
            height=400,
            showlegend=False
        )

        return fig

    def create_contribution_pie_chart(
        self,
        decomposition: MetricDecompositionResult,
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a pie chart showing component contributions.

        Args:
            decomposition: MetricDecompositionResult to visualize
            title: Optional chart title

        Returns:
            Plotly figure
        """
        # Prepare data
        components = list(decomposition.component_contributions.keys())
        contributions = [abs(decomposition.component_contributions[c]) * 100 for c in components]

        fig = go.Figure(data=[go.Pie(
            labels=components,
            values=contributions,
            textinfo='label+percent',
            hovertemplate='<b>%{label}</b><br>Contribution: %{value:.1f}%<extra></extra>'
        )])

        chart_title = title or f"Component Contributions - {decomposition.metric}"

        fig.update_layout(
            title=chart_title,
            template=self.default_template,
            height=400
        )

        return fig

    def create_anomaly_timeline(
        self,
        df: pd.DataFrame,
        metric_column: str,
        time_column: str,
        anomalies: List[AnomalyResult],
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a timeline chart highlighting anomalies.

        Args:
            df: DataFrame with time series data
            metric_column: Column name of the metric
            time_column: Column name for timestamps
            anomalies: List of detected anomalies
            title: Optional chart title

        Returns:
            Plotly figure
        """
        fig = go.Figure()

        # Plot main metric line
        fig.add_trace(go.Scatter(
            x=df[time_column],
            y=df[metric_column],
            mode='lines+markers',
            name='Metric',
            line=dict(color='steelblue', width=2),
            marker=dict(size=6)
        ))

        # Highlight anomalies
        if anomalies:
            anomaly_dates = [a.timestamp for a in anomalies]
            anomaly_values = [a.actual_value for a in anomalies]
            anomaly_colors = {
                'low': 'yellow',
                'medium': 'orange',
                'high': 'red',
                'critical': 'darkred'
            }
            colors = [anomaly_colors.get(a.severity.value, 'red') for a in anomalies]

            fig.add_trace(go.Scatter(
                x=anomaly_dates,
                y=anomaly_values,
                mode='markers',
                name='Anomalies',
                marker=dict(
                    size=12,
                    color=colors,
                    symbol='x',
                    line=dict(width=2, color='black')
                ),
                text=[f"Severity: {a.severity.value}<br>Deviation: {a.deviation_percentage:.1f}%"
                      for a in anomalies],
                hovertemplate='<b>Anomaly</b><br>Date: %{x}<br>Value: %{y:,.0f}<br>%{text}<extra></extra>'
            ))

        chart_title = title or f"Metric Timeline with Anomalies - {metric_column}"

        fig.update_layout(
            title=chart_title,
            xaxis_title="Date",
            yaxis_title=metric_column,
            template=self.default_template,
            height=400,
            hovermode='x unified'
        )

        return fig

    def create_comparison_chart(
        self,
        current_df: pd.DataFrame,
        baseline_df: pd.DataFrame,
        metric_column: str,
        dimension: str,
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a comparison chart between current and baseline periods.

        Args:
            current_df: Current period data
            baseline_df: Baseline period data
            metric_column: Column name of the metric
            dimension: Dimension to compare
            title: Optional chart title

        Returns:
            Plotly figure
        """
        # Aggregate by dimension
        current_agg = current_df.groupby(dimension)[metric_column].sum().reset_index()
        baseline_agg = baseline_df.groupby(dimension)[metric_column].sum().reset_index()

        # Merge
        merged = pd.merge(
            current_agg,
            baseline_agg,
            on=dimension,
            how='outer',
            suffixes=(' (Current)', ' (Baseline)')
        ).fillna(0)

        # Sort by current values
        merged = merged.sort_values(f"{metric_column} (Current)", ascending=False).head(10)

        fig = go.Figure()

        fig.add_trace(go.Bar(
            name='Baseline',
            x=merged[dimension],
            y=merged[f"{metric_column} (Baseline)"],
            marker_color='lightblue'
        ))

        fig.add_trace(go.Bar(
            name='Current',
            x=merged[dimension],
            y=merged[f"{metric_column} (Current)"],
            marker_color='steelblue'
        ))

        chart_title = title or f"Current vs Baseline - {metric_column} by {dimension}"

        fig.update_layout(
            title=chart_title,
            xaxis_title=dimension,
            yaxis_title=metric_column,
            barmode='group',
            template=self.default_template,
            height=400
        )

        return fig

    def create_multi_level_drill_down(
        self,
        drill_results: List[DrillDownResult],
        title: Optional[str] = None
    ) -> go.Figure:
        """
        Create a multi-level drill-down visualization.

        Args:
            drill_results: List of DrillDownResult for each level
            title: Optional chart title

        Returns:
            Plotly figure with subplots for each level
        """
        n_levels = len(drill_results)

        fig = make_subplots(
            rows=1,
            cols=n_levels,
            subplot_titles=[f"Level {i}: {r.dimension}" for i, r in enumerate(drill_results)]
        )

        for idx, result in enumerate(drill_results):
            if not result.top_contributors:
                continue

            labels = [c.value for c in result.top_contributors]
            values = [c.contribution_percentage for c in result.top_contributors]

            fig.add_trace(
                go.Bar(
                    x=labels,
                    y=values,
                    name=result.dimension,
                    marker_color='steelblue',
                    showlegend=False
                ),
                row=1,
                col=idx + 1
            )

        chart_title = title or "Multi-Level Drill-Down Analysis"

        fig.update_layout(
            title=chart_title,
            template=self.default_template,
            height=400
        )

        fig.update_yaxes(title_text="Contribution %", row=1, col=1)

        return fig

    def export_to_html(self, fig: go.Figure, filename: str):
        """Export figure to HTML file."""
        fig.write_html(filename)

    def export_to_image(self, fig: go.Figure, filename: str, format: str = 'png'):
        """Export figure to static image (requires kaleido)."""
        fig.write_image(filename, format=format)
